import json
from pydantic import BaseModel, Field
from langchain_core.messages import HumanMessage, AIMessage
# from langchain_deepseek import ChatDeepSeek
from langgraph.graph import StateGraph, MessagesState, END
from dotenv import load_dotenv

load_dotenv()

# 定义结构化输出模型
class AnalysisResult(BaseModel):
    """文本分析结果"""
    topic: str = Field(description="文本主要主题")
    summary: str = Field(description="内容摘要，不超过100字")
    keywords: list[str] = Field(description="提取的关键词列表")
    sentiment: str = Field(description="情感倾向：积极/消极/中性")

# 初始化模型
# model = ChatDeepSeek(model="deepseek-chat")
from langchain_ollama import ChatOllama
llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)
structured_llm = llm.with_structured_output(AnalysisResult).with_config(tags=["文本分析器"])

# 定义分析节点
def analysis_node(state: MessagesState):
    """对输入文本进行结构化分析"""
    result = structured_llm.invoke(state["messages"])
    print(f"分析结果: {result}")
    
    # 将分析结果转换为AIMessage
    response_content = f"""
📊 分析报告
主题: {result.topic}
摘要: {result.summary}
关键词: {', '.join(result.keywords)}
情感倾向: {result.sentiment}
    """.strip()
    
    return {"messages": [AIMessage(content=response_content)]}

# 构建LangGraph工作流
workflow = StateGraph(MessagesState)
workflow.add_node("分析", analysis_node)
workflow.set_entry_point("分析")
workflow.add_edge("分析", END)

# 编译图
graph = workflow.compile()

# 测试用例
def test_analysis():
    """测试文本分析功能"""
    test_text = """
    人工智能技术正在快速发展，特别是大语言模型的出现，
    为各个行业带来了前所未有的变革机会。从自动化客服到智能写作，
    从代码生成到数据分析，AI正在重塑我们的工作方式。
    虽然技术发展令人兴奋，但我们也需要关注AI伦理和安全性问题。
    """
    
    graph_input = {"messages": [HumanMessage(content=f"请分析以下文本：{test_text}")]}
    
    # 运行分析
    result = graph.invoke(graph_input)
    print("="*50)
    print(">>>最终输出:")
    # print(result)
    # print(json.dumps(result, indent=2, ensure_ascii=False))
    # print(result["messages"][-1].content)

# 执行测试
if __name__ == "__main__":
    test_analysis()